rm(list = ls())
library(tidyverse)
library(estimatr)
library(grid)
library(ri2)
library(reshape2)
library(bandit)
library(stargazer)
setwd(dirname(rstudioapi::getActiveDocumentContext()$path))
set.seed(38209)
source('fn.R')

load('data.rdata')
mod  <- lm_robust(Y ~ arm - 1, data = data, weights = ipw)
true <- coef(mod)
iter <- 10000

if (file.exists('Output/T2_simulation.rdata')){
  load('Output/T2_simulation.rdata')
}else {
  out_a <- simulate(periods = 10, n = 200, probs = true, iter = iter, static = FALSE, ppmat = TRUE)
  out_s <- simulate(periods = 10, n = 200, probs = true, iter = iter, static = TRUE, ppmat = TRUE)
}

adapt <- result_sum(data = as.data.frame(out_a$d_fit))
stati <- result_sum(data = as.data.frame(out_s$d_fit))
final <- bind_cols(adapt %>% dplyr::select(term, true, best, rmse, coverage), 
                   stati %>% dplyr::select(best, rmse, coverage)) 
names(final) <- c('Arm', 'True', 'Best A', 'RMSE A', 'Coverage A', 'Best', 'RMSE', 'Coverage')
stargazer::stargazer(final, rownames = F, summary = F, digits = 3)

mean(out_a$outmat[,'ses'])/mean(out_s$outmat[,'ses'])
